from querier.esquerier import ElasticSearchQuerier


class Word2VecInfoQuerier(ElasticSearchQuerier):

    def __init__(self, es, index, doc_type, nlp_service=None):
        super(Word2VecInfoQuerier, self).__init__(None, None, None)
        self.es = es
        self.index = index
        self.doc_type = doc_type
        self.nlp_service = nlp_service

    def search(self, args):
        year = args.get('year')
        week = args.get('week')
        word = args.get('word')
        id_ = word + year + week

        res = {}
        if id_ is not None:
            res = self.es.get(index=self.index, doc_type=self.doc_type, id=id_)
        return self._build_result(res, None)

    def _build_query(self, args):
        pass

    def _build_result(self, es_result, param):
        if es_result['found']:
            source_ = es_result['_source']

            similar_words = source_['similar_words']

            similar_words_weight = source_['similar_words_weight']

            return {
                'similar_words': similar_words,
                'similar_words_weight': similar_words_weight
            }
        else:
            return {}


class WordDfInfoQuerier(ElasticSearchQuerier):

    def __init__(self, es, index, doc_type, nlp_service=None):
        super(WordDfInfoQuerier, self).__init__(None, None, None)
        self.es = es
        self.index = index
        self.doc_type = doc_type
        self.nlp_service = nlp_service

    def search(self, args):
        year = args.get('year')
        week = args.get('week')
        word = args.get('word')
        id_ = word + year + week

        res = {}
        if id_ is not None:
            try:
                res = self.es.get(index=self.index, doc_type=self.doc_type, id=id_)
            except Exception as e:
                print(str(e))
                res = {}
        return self._build_result(res, None)

    def _build_query(self, args):
        pass

    def _build_result(self, es_result, param):
        if es_result.get('found'):
            source_ = es_result['_source']

            word = source_['word']

            df = source_['df']

            return {
                'word': word,
                'df': df
            }
        else:
            return {}


class Word2VecSearchQuerier(ElasticSearchQuerier):

    def __init__(self, word2vec_info, word_df_info):
        super(Word2VecSearchQuerier, self).__init__(None, None, None)
        self.word2vec_info = word2vec_info
        self.word_df_info = word_df_info

    def search(self, args):
        word = args.get('word')
        year = args.get('year')
        week = args.get('week')
        need_df = args.get('need_df')

        if need_df is None:
            need_df = True

        res = self.word2vec_info.search(args)

        similar_words = res.get('similar_words', [])
        similar_words_weight = res.get('similar_words_weight', [])

        df = []
        df_prev = []
        sk = similar_words
        skw = similar_words_weight

        prev_week = int(week) - 1

        if prev_week <= 0:
            prev_year = str(int(year) - 1)
            prev_week = '52'
        else:
            prev_week = str(prev_week).zfill(2)
            prev_year = year

        word_df = 0
        word_df_prev = 0

        if need_df:
            sk = []
            skw = []
            word_df_res = self.word_df_info.search({'word': word, 'year': year, 'week': week})
            word_df_res_prev = self.word_df_info.search({'word': word, 'year': prev_year, 'week': prev_week})

            if word_df_res and (not word_df_res.get('message')):
                word_df = word_df_res.get('df', 0)
                word_df_prev = word_df_res_prev.get('df', 0)

            for k, w in zip(similar_words, similar_words_weight):
                df_res = self.word_df_info.search({'word': k, 'year': year, 'week': week})
                df_res_prev = self.word_df_info.search({'word': k, 'year': prev_year, 'week': prev_week})
                if df_res and (not df_res.get('message')):
                    df.append(df_res.get('df', 0))
                    df_prev.append(df_res_prev.get('df', 0))
                    sk.append(k)
                    skw.append(w)

        # if not word_df:
        #     return {}

        out = {
            'word': word,
            'word_df': word_df,
            'word_df_change': word_df - word_df_prev,
            'word_weight': 1.0,
            'similar_words': sk,
            'similar_words_weight': skw,
            'similar_words_df': df,
            'similar_words_df_change': [a - b for a, b in zip(df, df_prev)]
        }
        return self._build_result(out, None)

    def _build_query(self, args):
        pass

    def _build_result(self, es_result, param):
        return es_result
